package com.atao.baokede.rag;

import org.springframework.ai.document.Document;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

@Component
public class DocumentLoader {
    @Autowired
    private ResourcePatternResolver resourcePatternResolver;

    /**
     * 加载所有markdown文件
     * @return 文档列表
     */
    public List<Document> loadMarkdowns() throws IOException {
        List<Document> documents = new ArrayList<>();
        Resource[] resources = resourcePatternResolver.getResources("classpath:document/*.md");
        for (Resource resource : resources){
            String filename = resource.getFilename();
            MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                    .withHorizontalRuleCreateDocument(true)
                    .withIncludeCodeBlock(false)
                    .withIncludeBlockquote(false)
                    .withAdditionalMetadata("filename", filename)
                    .build();
            MarkdownDocumentReader reader = new MarkdownDocumentReader(resource,config);
            documents.addAll(reader.get());
        }
        return documents;
    }
}
